package com.xuhui.xuaiagent.agent;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.thread.ThreadUtil;
import cn.hutool.core.util.StrUtil;
import com.xuhui.xuaiagent.agent.model.AgentState;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;


import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * 抽象基础代理类 用于管理代理状态和执行流程
 *
 * 提供状态转换 内存管理和基于步骤的执行循环的基础功能
 * 子类必须实现step方法
 * @author 21829
 */
@Data
@Slf4j
public abstract class BaseAgent {
    // 核心属性
    private String name;

    // 提示词
    private String systemPrompt;
    private String nextStepPrompt;

    // 代理状态
    private AgentState state = AgentState.IDLE;

    // 执行控制
    private int MaxSteps = 10;
    private int currentStep = 0;

    // LLM 大模型
    private ChatClient chatClient;

    // memory (需要自主维护会话上下文)
    private List<Message> messageList = new ArrayList<>();
    // 存储ai思考信息和最终响应
    private List<String> AiResponceList = new ArrayList<>();


    /**
     * 运行代理
     *
     * @param userPrompt 用户提示词
     * @return 执行结果
     */
    public String run(String userPrompt){
        //1.1 判断当前运行状态是否为空闲状态
        if (this.state != AgentState.IDLE){
            throw new RuntimeException("Cannot run agent form state:" + this.state);
        }
        //1.2 判断用户提示词是否为空
        if (StrUtil.isBlank(userPrompt)){
            throw new RuntimeException("Cannot run agent with empty user prompt");
        }
        //2 如果上面判断没问题 可以运行 更改状态
        state = AgentState.RUNNING;
        //3 记录消息上下文
        messageList.add(new UserMessage(userPrompt));
        // 创建一个用来保存执行结果的列表
        List<String> results = new ArrayList<>();
        // 开始循环执行
        try {
            for (int i = 0; i < MaxSteps && state != AgentState.FINISHED; i++){
                int stepNumber = i + 1;
                currentStep = stepNumber;
                log.info("Executing step" + stepNumber + "/" + MaxSteps);

                // 单步执行
                String stepResult = step();
                String result = "Step" + stepNumber + ":" + stepResult;
                results.add(result);
            }
            // 检查是否超出步骤限制
            if (currentStep >= MaxSteps){
                state = AgentState.FINISHED;
                results.add("Terminated: Reached max steps (" + MaxSteps + ")");
            }
            return String.join("\n",results);
        } catch (Exception e) {
            state = AgentState.ERROR;
            log.error("Error executing agent", e);
            return "执行错误" + e.getMessage();
        } finally {
            this.cleanup();
        }
    }

    /**
     * 执行单个步骤
     * @return 步骤执行结果
     */
    public abstract  String step();

    /**
     * 清理资源
     */
    protected void cleanup(){
        // 子类可以通过重写此方法来清理资源
    }

    /**
     * 运行代理（流式输出）
     *
     * @param userPrompt 用户提示词
     * @return SseEmitter实例
     */
    public SseEmitter runStream(String userPrompt){
        // 创建SseEmitter。设置较长的超时时间
        SseEmitter emitter = new SseEmitter(3000000L);

        // 使用线程异步处理 避免阻塞主线程
        CompletableFuture.runAsync(() -> {
          try { // 如果当前状态不等于空闲状态
              if (this.state != AgentState.IDLE){
                  emitter.send("错误: 无法从状态运行代理:" + this.state);
                  emitter.complete();
                  return;
              }
              log.error("异步任务执行了");
              if (StrUtil.isBlank(userPrompt)){
                  emitter.send("错误: 不能使用空提示词运行代理");
                  emitter.complete();
                  return;
              }
              // 更改状态
              state = AgentState.RUNNING;
              // 记录消息上下文
              messageList.add(new UserMessage(userPrompt));

              try {
                  for (int i = 0; i < MaxSteps && state != AgentState.FINISHED; i++){
                      int stepNumber = i + 1;
                      currentStep = stepNumber;
                      log.info("Executing step" + stepNumber + "/" + MaxSteps);

                      // 单步执行
                      String stepResult = step();
                      String result = "Step" + stepNumber + ":" + stepResult;
                      if (CollUtil.isNotEmpty(AiResponceList)){
                          String thinking = getAiResponceList().get(i);
                          emitter.send(thinking);
                      }
                      // 发送每一步的结果
                      log.error(result);
//                      emitter.send(result);
                  }
                  // 检查是否超出步骤范围
                  if (currentStep >= MaxSteps){
                      state = AgentState.FINISHED;
                      emitter.send("执行结束: 达到最大步骤 (" + MaxSteps + ")");
                  }
                  // 执行到此时则代表是正常完成
                  emitter.complete();
              }catch (Exception e){
                  state = AgentState.ERROR;
                  log.error("执行智能体失败",e);
                  try {
                      emitter.send("执行错误:" + e.getMessage());
                      emitter.complete();
                  }catch (Exception ex){
                      emitter.completeWithError(ex);
                  }
              }finally {
                  // 清理资源
                  this.cleanup();
              }
          }catch (Exception e){
              emitter.completeWithError(e);
          }
        });
        // 设置超时时间和完成回调
        emitter.onTimeout(() -> {
            this.state = AgentState.ERROR;
            this.cleanup();
            log.warn("SSE connection timed out");
        });
        emitter.onCompletion(() -> {
            if (state == AgentState.RUNNING){
                state = AgentState.FINISHED;
            }
            this.cleanup();
            log.info("SSE connection terminated");
        });
        return emitter;

    }


}
